#ifndef TENSOR_H
#define TENSOR_H

#include <iostream>
#include <memory>
#include <string>
#include <functional>
#include <unordered_set>
#include <vector>
#include <algorithm> // 用于 reverse 函数
#include <cmath>
#include <cstdlib>
#include <random> // 添加这一行以使用随机数生成器
#include <stdexcept>
#include <iomanip>
using namespace std;

class Value
{
public:
    // 重载输出运算符
    friend ostream &operator<<(ostream &os, const Value &v);
    friend Value operator+(double val, Value &v);
    friend Value operator*(double val, Value &v);

public:
    Value() : data(0.0), label(""), _op(""), grad(0.0), _backward([]() {}), _prev({}) {}
    Value(double d, string label = "", string _op = "", unordered_set<Value *> _prev = {}) : data(d), label(label), _op(_op), grad(0.0), _backward([]() {}), _prev(_prev) {}
    Value(const Value &v) : data(v.data), label(v.label), _op(v._op), grad(v.grad), _backward(v._backward), _prev(v._prev) {}

    /* docs[1] Thinking for this
       for x in xs:
          sum = sum + x

       if we reset the sum to new_sum, the old_sum 's point is reset too(wrong)
       that's to say, we shoule let the new_sum be a new obj not change the old one.
       so we should make a new obj and return it like:

       A* sum;
       vector<A*> xs;
       for x in xs:
          sum = new A(*sum + *x)

       but how can we release the old_sum?
       DO NOT forget we have the _prev(vector<A*>_children), we can use it.

       SO if we need a not-fixed Value(such as loss, sum...), we must use A* non_fixed = new A().
    */
    Value &operator=(const Value &v)
    {
        if (this != &v)
        {
            data = v.data;
            label = v.label;
            _op = v._op;
            grad = v.grad;
            _backward = v._backward;
            _prev = v._prev;
        }
        return *this;
    }

    //  Value op Value
    Value operator+(Value &rhs)
    {
        Value out(data + rhs.data, "", "+", {const_cast<Value *>(this), &rhs});
        out._backward = [this, &rhs, &out]() mutable
        {
            const_cast<Value *>(this)->grad += 1.0 * out.grad;
            rhs.grad += 1.0 * out.grad;
        };
        return out;
    }

    Value operator*(Value &rhs) const
    {
        Value out(data * rhs.data, "", "*", {const_cast<Value *>(this), &rhs});
        out._backward = [this, &rhs, &out]() mutable
        {
            const_cast<Value *>(this)->grad += rhs.data * out.grad;
            rhs.grad += this->data * out.grad;
        };
        return out;
    }

    Value relu() const
    {
        Value out(data > 0 ? data : 0, "", "relu", {const_cast<Value *>(this)});
        out._backward = [this, &out]() mutable
        {
            const_cast<Value *>(this)->grad += (data > 0) * out.grad;
        };
        return out;
    }

    Value tanh() const
    {
        Value out(std::tanh(data), "", "tanh", {const_cast<Value *>(this)});
        out._backward = [this, &out]() mutable
        {
            const_cast<Value *>(this)->grad += (1 - out.data) * out.grad;
        };
        return out;
    }
    Value operator-(Value &rhs)
    {
        Value *tmp = new Value(rhs * -1);
        return (*this) + *tmp;
    }
    Value operator-()
    {
        return *this * -1;
    }

    // Value op double
    Value operator+(double rhs_value)
    {
        Value *tmp = new Value(rhs_value);
        return (*this) + *tmp; // TODO tmp 是局部变量会被销毁, 所以只能new，但是可能内存泄露,所以靠out._prev释放的时候释放
    }
    Value operator*(double rhs_value)
    {
        Value *tmp = new Value(rhs_value);
        return (*this) * *tmp;
    }
    Value operator-(double rhs_value)
    {
        Value *tmp = new Value(rhs_value);
        return (*this) - *tmp;
    }

    void backward(bool verbose = false)
    {
        // topo sort
        std::unordered_set<Value *> visited;
        std::vector<Value *> topo;

        std::function<void(Value *)> build_topo = [&](Value *v)
        {
            if (visited.find(v) == visited.end())
            {
                visited.insert(v);
                for (auto &n : v->_prev)
                {
                    build_topo(n);
                }
                topo.push_back(v);
            }
        };
        build_topo(this);
        // cout << "topo size:" << topo.size() << endl;

        grad = 1.0;

        for (auto begin = topo.rbegin(); begin != topo.rend(); ++begin)
        {
            if (verbose)
                cout << "topo sequence:" << (*begin)->label << "and it'grad:" << (*begin)->grad << endl;
            if ((*begin)->_backward != nullptr)
                (*begin)->_backward();
        }
    }

    ~Value() = default;

    double data;
    double grad;
    string label;
    string _op;
    function<void()> _backward;
    unordered_set<Value *> _prev;
};

ostream &operator<<(ostream &os, const Value &v)
{
    char buffer[1024];
    sprintf(buffer, "Value(data=%.8f, label=%s, grad=%.8f)", v.data, v.label.c_str(), v.grad);
    cout << buffer;
    return os;
}

Value operator+(double val, Value &v)
{
    return v + val;
}
Value operator*(double val, Value &v)
{
    return v * val;
}

class Neuron
{
public:
    Neuron(int n_in, bool bias = true, string label = "") : W(n_in), bias(bias), label(label)
    {
        // 使用标准正态分布初始化权重
        std::random_device rd;
        std::mt19937 gen(rd());
        std::normal_distribution<> dist(0.0, 1.0); // 均值为0，标准差为1

        for (int i = 0; i < n_in; i++)
        {
            W[i] = new Value(dist(gen), label + "w" + to_string(i), "w");
        }
        b = new Value(dist(gen), label + "b", "b");
    }

    Value *operator()(vector<Value *> &x)
    {
        // TODO 含有大量中间对象，但是不能delete 他们是无名的英雄！
        Value *out = new Value(0.0);
        for (int i = 0; i < x.size(); i++)
        {
            Value *tmp = new Value(*x[i] * *W[i]);
            out = new Value(*tmp + *out);
        }
        out = new Value(out->tanh());

        if (bias)
        {
            Value *ans = new Value(*out + *b);
            return ans;
        }
        else
        {
            Value *ans = new Value(*out);
            return ans;
        }
    }

    vector<Value *> parameters()
    {
        vector<Value *> out(W.begin(), W.end());
        if (bias)
            out.push_back(b);
        return out;
    }

    vector<Value *> W;
    Value *b;
    bool bias;
    string label;

private:
};

class Layer
{
public:
    Layer(int n_in, int n_out, bool bias = true, string label = "") : neurons(n_out)
    {
        srand(static_cast<unsigned int>(time(nullptr)));
        for (int i = 0; i < n_out; i++)
        {
            neurons[i] = new Neuron(n_in, bias, label + "neuron" + to_string(i) + "-");
        }
    }

    vector<Value *> operator()(vector<Value *> &x)
    {
        vector<Value *> out;
        for (int i = 0; i < neurons.size(); i++)
        {
            out.push_back((*neurons[i])(x));
        }
        return out;
    }

    vector<Value *> parameters()
    {
        vector<Value *> out;
        for (int i = 0; i < neurons.size(); i++)
        {
            vector<Value *> tmp = neurons[i]->parameters();
            out.insert(out.end(), tmp.begin(), tmp.end());
        }
        return out;
    }

private:
    vector<Neuron *> neurons;
    string label;
};

class Tensor
{
public:
    friend ostream &operator<<(ostream &os, const Tensor &t);

public:
    // 构造函数：随机初始化 for data
    Tensor(const std::vector<int> &dimensions, double min_val = 0.0, double max_val = 1.0)
    {
        initialize(dimensions, min_val, max_val);
    }

    // 构造函数：从 vector 初始化 for data
    Tensor(const std::vector<double> &data, const std::vector<int> &dimensions = {1})
    {
        initialize(data, dimensions);
    }
    Tensor(const vector<Value *> &data, const std::vector<int> &dimensions)
    {
        _dimensions = dimensions;
        _data = data;
    }

    Tensor(const Tensor &&tensor)
    {
        _dimensions = std::move(tensor._dimensions);
        _data = std::move(tensor._data);
    }

    Tensor(const Tensor &tensor) : _dimensions(tensor._dimensions), _data(tensor._data)
    {
        // _data = tensor._data; // TODO(mx) share assignment?
    }

    // 获取维度
    std::vector<int> getDimensions() const
    {
        return _dimensions;
    }

    // 获取数据
    const std::vector<Value *> &getData() const
    {
        return _data;
    }

    // 打印 Tensor
    void print() const
    {
        std::cout << "Tensor with dimensions: ";
        for (int dim : _dimensions)
        {
            std::cout << dim << " ";
        }
        std::cout << "\nData: ";
        for (const auto &val : _data)
        {
            std::cout << val->data << " ";
        }
        std::cout << std::endl;
    }

    // 多维索引访问
    double at(const std::vector<int> &indices) const
    {
        int index = 0;
        int stride = 1;
        for (int i = _dimensions.size() - 1; i >= 0; --i)
        {
            index += indices[i] * stride;
            stride *= _dimensions[i];
        }
        return _data[index]->data;
    }

    double operator[](const std::vector<int> &indices) const
    {
        int index = 0;
        int stride = 1;
        for (int i = _dimensions.size() - 1; i >= 0; --i)
        {
            index += indices[i] * stride;
            stride *= _dimensions[i];
        }
        return _data[index]->data;
    }

    /* math functions
     * What we should do is to make the result once and return it wrappered by Tensor.
     */

    // +
    Tensor operator+(const Tensor &other) const
    {
        _check_dim(other._dimensions.size() - 1, _dimensions);
        vector<Value *> ans(_data.size());
        for (size_t i = 0; i < _data.size(); ++i)
            ans[i] = new Value(*_data[i] + *other._data[i]);
        return Tensor(ans, _dimensions);
    }

    Tensor operator*(const Tensor &other) const
    {
        _check_dim(other._dimensions.size() - 1, _dimensions);
        vector<Value *> ans(_data.size());
        for (size_t i = 0; i < _data.size(); ++i)
            ans[i] = new Value(*_data[i] * *other._data[i]);
        return Tensor(ans, _dimensions);
    }

    // sum
    Tensor sum(int dim = -1, bool keep_dims = false)
    {
        _check_dim(dim, _dimensions);
        Value *sum = new Value(0.0);

        for (auto &val : _data)
        {
            sum = new Value(*sum + *val);
        }
        return Tensor(vector<Value *>({sum}), {1});
    }

    // 改变视图
    Tensor view(const std::vector<int> &new_dimensions)
    {
        size_t sz = _data.size();
        _check_sz(sz, new_dimensions);
        return Tensor(_data, new_dimensions);
    }

    void backward(bool verbose = false)
    {
        _check_sz(1, _dimensions);
        _data[0]->backward(verbose);
    }

private:
    std::vector<Value *> _data;
    std::vector<int> _dimensions;

    // 随机初始化
    void initialize(const std::vector<int> &dimensions, double min_val, double max_val)
    {
        _dimensions = dimensions;
        size_t sz = accumulate(dimensions.begin(), dimensions.end(), 1, std::multiplies<int>());
        _data.resize(sz);

        std::random_device rd;
        std::mt19937 gen(rd());
        std::uniform_real_distribution<> dis(min_val, max_val);

        for (int i = 0; i < sz; ++i)
        {
            _data[i] = new Value(dis(gen));
        }
    }

    void _check_sz(size_t sz, const std::vector<int> &dimensions) const
    {
        if (sz != accumulate(dimensions.begin(), dimensions.end(), 1, std::multiplies<int>()))
            throw std::invalid_argument("Data size does not match the specified dimensions");
    }

    void _check_dim(int sz, const std::vector<int> &dimensions) const
    {
        int sz_dim = dimensions.size();
        if (sz >= sz_dim)
        {
            throw std::invalid_argument("Invalid dimension index.");
        }
    }

    // 从 vector 初始化
    void initialize(const std::vector<double> &data, const std::vector<int> &dimensions)
    {
        size_t sz = data.size();
        _check_sz(sz, dimensions);
        _data.resize(sz);
        for (int i = 0; i < sz; ++i)
        {
            _data[i] = new Value(data[i]);
        }
        _dimensions = dimensions;
    }
};

// 重载 << 操作符
std::ostream &operator<<(std::ostream &os, const Tensor &t)
{
    const std::vector<int> &dimensions = t.getDimensions();
    const std::vector<Value *> &data = t.getData();

    // 计算 stride 的局部函数
    auto stride = [&dimensions](int depth) -> int
    {
        int s = 1;
        for (int i = depth; i < dimensions.size(); ++i)
        {
            s *= dimensions[i];
        }
        return s;
    };

    // 递归打印函数
    auto print = [&os, &data, &dimensions, &stride](auto &&print, int depth, int index) -> void
    {
        if (depth == dimensions.size())
        {
            if (index >= data.size())
            {
                return;
            }
            os << std::setw(3) << data[index]->data << ",";
            return;
        }

        os << "[";
        for (int i = 0; i < dimensions[depth]; ++i)
        {
            print(print, depth + 1, index + i * stride(depth + 1));
            if (i < dimensions[depth] - 1)
            {
                os << "\n"
                   << std::string(depth * 2, ' ');
            }
        }
        os << "]";
        if (depth > 0)
        {
            os << " ";
        }
    };

    // 开始打印
    os << "Tensor(";
    os << setprecision(6);
    print(print, 0, 0);
    // 添加维度信息
    os << ", shapes=(";
    for (int i = 0; i < dimensions.size(); ++i)
    {
        i &&os << ",";
        os << dimensions[i];
    }
    if (dimensions.size() == 1)
        os << ",";
    os << "))";
    os << defaultfloat;
    return os;
}
#endif // TENSOR_H